# Simple Firewall
from collections import defaultdict
from time import time_ns

from scapy.compat import raw
from scapy.layers.inet import IP, TCP, UDP

from NFs.firewall_rule import FirewallRule


class Firewall(object):
    def __init__(self, csv_file='fw_rules.csv'):
        self.port_to_rules = defaultdict(FirewallRule)  # 规则集
        self.csv_file = csv_file

    def parse_rules(self):  #
        with open(self.csv_file, 'r') as f:
            for line in f:
                direction, protocol, port, ip = line.strip().split('\t')
                # 规则中端口是范围
                if '-' in port:
                    lo, hi = port.strip().split('-')
                    for i in range(int(lo), int(hi) + 1):
                        firewall_rule = self.port_to_rules[i]  # 默认返回空的FirewallRule
                        firewall_rule.add_direction(direction)
                        firewall_rule.add_protocol(protocol)
                        if '-' in ip:
                            firewall_rule.add_ips_range(ip)
                        else:
                            firewall_rule.add_ips_discrete(ip)
                else:
                    firewall_rule = self.port_to_rules[int(port)]
                    firewall_rule.add_direction(direction)
                    firewall_rule.add_protocol(protocol)
                    if '-' in ip:
                        firewall_rule.add_ips_range(ip)
                    else:
                        firewall_rule.add_ips_discrete(ip)

    def accept_packet(self, direction: str, protocol: str, port: int, ip_address: str) -> bool:
        rule = self.port_to_rules[port]
        return rule.contains_direction(direction) and \
               rule.contains_protocol(protocol) and rule.is_within_ip_range(ip_address)

    def match(self, pkt: bytes) -> bool:
        scapy_pkt = IP(pkt)
        proto = scapy_pkt[IP].proto
        if proto == 6:
            protocol = 'tcp'
        else:  # 17-- UDP
            protocol = 'udp'
        src, dst = scapy_pkt[IP].src, scapy_pkt[IP].dst
        if proto == 6:  # tcp
            sport, dport = scapy_pkt[TCP].sport, scapy_pkt[TCP].dport
        else:  # 17
            sport, dport = scapy_pkt[UDP].sport, scapy_pkt[UDP].dport
        return not (self.accept_packet('inbound', protocol, sport, src) or self.accept_packet('outbind', protocol,
                                                                                              dport, dst))


if __name__ == '__main__':
    start = time_ns()
    firewall = Firewall()
    firewall.parse_rules()
    print(firewall.port_to_rules[19953].ips_discrete, )
    print((time_ns() - start) / 10 ** 9)
    pkt = b'E\x00\x00@\x00\x01\x00\x00@\x06\x11\xe1\xf7\x89\xe9\xcc$\x91b\xeffeM\xf1\x00\x00\x00\x00\x00\x00\x00\x00P' \
          b'\x02 \x00\xe2\x0c\x00\x00aaaaaaaaaaaaaaaaaaaaaaaa '

    pkt = IP(src='247.137.233.204', dst='36.145.98.239') / TCP(sport=26213, dport=19953)
    pkt[TCP].add_payload(b'a' * 3966)
    pkt = raw(pkt)
    count = 1
    start = time_ns()
    for i in range(100000):
        flag = firewall.match(pkt)
    print((time_ns() - start) / 100000 / 1000)  # 平均每条指令执行时间
